import glob
import re
import random
import os

import numpy as np
import PIL
import skimage.util

import mat_util
import image_util
from automap import train_util

ARTEFACT = 'artifact'
NOISE = 'noise'
DETAIL_LESS = 'detail_less'


# 将指定文件夹的图像(默认png格式)转成降采样数据和原数据,并批量保存为一个mat
def get_trained_data(data_path, data_suffix='.png', save_suffix='.png', shape=(256, 256), type=ARTEFACT,
                     count_per_data=1000, gausse_max=0.35, salt_amount_max=0.35,
                     max_artifact_interval=5, dowsample_radius_ratio=(0.08, 0.2)):
    print(data_path)
    print(type)
    # 初始化目标文件夹开始
    downsample_outpath = data_path + '\downsample'
    if not os.path.exists(downsample_outpath):
        os.makedirs(downsample_outpath)

    train_data_path = data_path + '\\train'
    if not os.path.exists(train_data_path):
        os.makedirs(train_data_path)
    if type == ARTEFACT:
        artifact_downsample_data_path = downsample_outpath + '\\' + ARTEFACT
        if not os.path.exists(artifact_downsample_data_path):
            os.makedirs(artifact_downsample_data_path)
        artifact_train_data_path = train_data_path + '\\' + ARTEFACT
        if not os.path.exists(artifact_train_data_path):
            os.makedirs(artifact_train_data_path)
        downsample_outpath = artifact_downsample_data_path
        train_data_path = artifact_train_data_path
    elif type == NOISE:
        noise_downsample_data_path = downsample_outpath + '\\' + NOISE
        if not os.path.exists(noise_downsample_data_path):
            os.makedirs(noise_downsample_data_path)
        noise_train_data_path = train_data_path + '\\' + NOISE
        if not os.path.exists(noise_train_data_path):
            os.makedirs(noise_train_data_path)
        downsample_outpath = noise_downsample_data_path
        train_data_path = noise_train_data_path
    elif type == DETAIL_LESS:
        dowsample_downsample_data_path = downsample_outpath + '\\' + DETAIL_LESS
        if not os.path.exists(dowsample_downsample_data_path):
            os.makedirs(dowsample_downsample_data_path)
        dowsample_train_data_path = train_data_path + '\\' + DETAIL_LESS
        if not os.path.exists(dowsample_train_data_path):
            os.makedirs(dowsample_train_data_path)
        downsample_outpath = dowsample_downsample_data_path
        train_data_path = dowsample_train_data_path
    # 初始化文件夹结束

    # 拼接np数组保存为npz文件
    x = np.zeros((count_per_data, shape[0], shape[1]), dtype=np.uint8)
    y = np.zeros((count_per_data, shape[0], shape[1]), dtype=np.uint8)

    listdir = glob.glob(data_path + r"\*" + data_suffix)
    listdir.sort(key=lambda name: int(re.findall('\d+', name)[0]))
    l = len(listdir)
    left_img = l % count_per_data  # 剩余的作为测试数据
    i = 0
    for file in listdir:
        if not file.endswith(data_suffix):
            continue
        cur_img = PIL.Image.open(file)  # 图像数据

        # plt.subplot(331)
        # plt.imshow(cur_img, cmap="gray")
        # plt.title("1")

        if not (shape[0] == cur_img.size[0] and shape[1] == cur_img.size[1]):
            cur_img = cur_img.resize(shape, PIL.Image.Resampling.LANCZOS)

            # plt.subplot(332)
            # plt.imshow(cur_img, cmap="gray")
            # plt.title("2")

        cur_img = np.array(cur_img)

        # plt.subplot(333)
        # plt.imshow(cur_img, cmap="gray")
        # plt.title("3")

        y[i % count_per_data] = cur_img

        # plt.subplot(334)
        # plt.imshow(y[i%count_per_data], cmap="gray")
        # plt.title("4")

        # 降采样处理
        # 只添加伪影
        if type == ARTEFACT:
            downsample_data = image_util.add_artifact(cur_img, random.randint(0, 2), random.randint(0, 2),
                                                      random.randint(2, max_artifact_interval + 1),
                                                      random.uniform(0.005, 0.01))
        # 只添加噪声
        elif type == NOISE:
            salt_amount = random.uniform(0.01, salt_amount_max)
            a = i % 4
            if a == 0:  # 加高斯噪声
                mean = random.uniform(0.01, gausse_max)
                var = random.uniform(0.01, gausse_max - mean)
                downsample_data = skimage.util.random_noise(cur_img, mode="gaussian",
                                                            var=var,
                                                            mean=mean)
            elif a == 1:
                downsample_data = skimage.util.random_noise(cur_img, mode='salt',
                                                            amount=salt_amount)
            elif a == 2:
                downsample_data = skimage.util.random_noise(cur_img, mode='pepper',
                                                            amount=salt_amount)
            elif a == 3:
                downsample_data = skimage.util.random_noise(cur_img, mode='s&p',
                                                            amount=salt_amount)
            downsample_data = downsample_data * 255
        elif type == DETAIL_LESS:
            # 圆型掩膜
            downsample_data = image_util.circle_downsample(cur_img, random.uniform(dowsample_radius_ratio[0],
                                                                                   dowsample_radius_ratio[1]))

        # plt.subplot(335)
        # plt.imshow(downsample_data, cmap="gray")
        # plt.title("6")

        x[i % count_per_data] = downsample_data.astype(int)
        # 重新转换
        downsample_image = PIL.Image.fromarray(downsample_data).convert("L")

        # plt.subplot(336)
        # plt.imshow(x[i%count_per_data], cmap="gray")
        # plt.title("6")
        # plt.show()

        new_name = '\\' + str(i)
        downsample_image.save(downsample_outpath + new_name + save_suffix)
        # imageio.imwrite(downsample_outpath + new_name + save_suffix, downsample_image)#用于肉眼对照

        i = i + 1
        if i % count_per_data == 0:
            print(i)
            index = i // count_per_data
            x_name = '{}_x{}.mat'.format(type, index)
            y_name = x_name.replace('x', 'y')
            x_file = os.path.join(train_data_path, x_name)
            y_file = os.path.join(train_data_path, y_name)
            mat_util.save_mat(x_file, x)
            mat_util.save_mat(y_file, y)
            print('{}和{}保存至:{}'.format(x_name, y_name, train_data_path))

    x = x[0:left_img - 1]
    y = y[0:left_img - 1]
    mat_util.save_mat(os.path.join(train_data_path, '{}_test.mat'.format(type)), y)
    mat_util.save_mat(os.path.join(train_data_path, '{}_origin.mat'.format(type)), x)


if __name__ == '__main__':
    path = r'E:\download\dataset\keras\IXI-T1\output'
    get_trained_data(path, type=ARTEFACT)
    get_trained_data(path, type=NOISE)
    get_trained_data(path, type=DETAIL_LESS)

    path = (r'E:\download\dataset\keras\IXI-T2\output')
    get_trained_data(path, type=ARTEFACT)
    get_trained_data(path, type=NOISE)
    get_trained_data(path, type=DETAIL_LESS)

    path = r'E:\download\dataset\keras\IXI-PD\output'
    get_trained_data(path, type=ARTEFACT)
    get_trained_data(path, type=NOISE)
    get_trained_data(path, type=DETAIL_LESS)

    path = r'E:\download\dataset\keras\IXI-MRA\output'
    get_trained_data(path, type=ARTEFACT)
    get_trained_data(path, type=NOISE, gausse_max=0.2, salt_amount_max=0.2)
    get_trained_data(path, type=DETAIL_LESS)

    path = r'E:\download\dataset\keras\IXI-DTI\output'
    get_trained_data(path, type=ARTEFACT)
    get_trained_data(path, type=NOISE)
    get_trained_data(path, type=DETAIL_LESS)

    # 拼接训练数据
    mat_util.joint_xy((r'E:\download\Dataset\keras\IXI-T1\output\train\artifact'
              , r'E:\download\Dataset\keras\IXI-T2\output\train\artifact'
              , r'E:\download\Dataset\keras\IXI-PD\output\train\artifact'
              , r'E:\download\Dataset\keras\IXI-DTI\output\train\artifact'
              , r'E:\download\Dataset\keras\IXI-MRA\output\train\artifact'),
             save_path=r'E:\download\Dataset\keras\train\local\artifact', per_mat=4, total_mat=4, start_data=1)

    mat_util.joint_xy((r'E:\download\Dataset\keras\IXI-T1\output\train\noise'
              , r'E:\download\Dataset\keras\IXI-T2\output\train\noise'
              , r'E:\download\Dataset\keras\IXI-PD\output\train\noise'
              , r'E:\download\Dataset\keras\IXI-DTI\output\train\noise'
              , r'E:\download\Dataset\keras\IXI-MRA\output\train\noise'),
             save_path=r'E:\download\Dataset\keras\train\local\noise', per_mat=4, total_mat=4, start_data=1)

    mat_util.joint_xy((r'E:\download\Dataset\keras\IXI-T1\output\train\detail_less'
              , r'E:\download\Dataset\keras\IXI-T2\output\train\detail_less'
              , r'E:\download\Dataset\keras\IXI-PD\output\train\detail_less'
              , r'E:\download\Dataset\keras\IXI-DTI\output\train\detail_less'
              , r'E:\download\Dataset\keras\IXI-MRA\output\train\detail_less'),
             save_path=r'E:\download\Dataset\keras\train\local\detail_less', per_mat=4, total_mat=4, start_data=1)

    # 拼接测试数据
    mat_util.joint_tests((r'E:\download\Dataset\keras\IXI-T1\output\train\detail_less'
                          , r'E:\download\Dataset\keras\IXI-T2\output\train\detail_less'
                          , r'E:\download\Dataset\keras\IXI-PD\output\train\detail_less'
                          , r'E:\download\Dataset\keras\IXI-DTI\output\train\detail_less'
                          , r'E:\download\Dataset\keras\IXI-MRA\output\train\detail_less'),
                         save_path=r'E:\download\Dataset\keras\train\local\detail_less')
    mat_util.joint_tests((r'E:\download\Dataset\keras\IXI-T1\output\train\noise'
                          , r'E:\download\Dataset\keras\IXI-T2\output\train\noise'
                          , r'E:\download\Dataset\keras\IXI-PD\output\train\noise'
                          , r'E:\download\Dataset\keras\IXI-DTI\output\train\noise'
                          , r'E:\download\Dataset\keras\IXI-MRA\output\train\noise'),
                         save_path=r'E:\download\Dataset\keras\train\local\noise')
    mat_util.joint_tests((r'E:\download\Dataset\keras\IXI-T1\output\train\artifact'
                          , r'E:\download\Dataset\keras\IXI-T2\output\train\artifact'
                          , r'E:\download\Dataset\keras\IXI-PD\output\train\artifact'
                          , r'E:\download\Dataset\keras\IXI-DTI\output\train\artifact'
                          , r'E:\download\Dataset\keras\IXI-MRA\output\train\artifact'),
                         save_path=r'E:\download\Dataset\keras\train\local\artifact')
    # train_path = r'E:\download\Dataset\keras\train\local\noise'
    # train_util.train_mat(train_path)
    # train_path = r'E:\download\Dataset\keras\train\local\artifact'
    # train_util.train_mat(train_path)
    # train_path = r'E:\download\Dataset\keras\train\local\detail_less'
    # train_util.train_mat(train_path)
